import numpy as np
import vtk
from vtk.util import numpy_support

def vtk_to_nas(vtk_file, nas_file):
    """
    将VTK 2.0格式的网格文件转换为NASTRAN格式
    
    参数:
    vtk_file: 输入的VTK文件路径
    nas_file: 输出的NASTRAN文件路径
    """
    # 1. 读取VTK文件
    reader = vtk.vtkUnstructuredGridReader()
    reader.SetFileName(vtk_file)
    reader.Update()
    grid = reader.GetOutput()
    
    # 2. 获取点坐标
    points = numpy_support.vtk_to_numpy(grid.GetPoints().GetData())
    
    # 3. 获取单元信息
    cells = []
    cell_types = []
    for i in range(grid.GetNumberOfCells()):
        cell = grid.GetCell(i)
        cell_points = []
        for j in range(cell.GetNumberOfPoints()):
            cell_points.append(cell.GetPointId(j))
        cells.append(cell_points)
        cell_types.append(cell.GetCellType())
    cells = np.array(cells)
    
    # 4. 写入NASTRAN文件
    with open(nas_file, 'w') as f:
        # 写入文件头
        f.write("$ NASTRAN Input Deck Generated by VTK converter\n")
        f.write("SOL 101\n")  # 添加求解类型
        f.write("CEND\n")
        f.write("BEGIN BULK\n")
        
        # 添加坐标系统定义
        f.write("CORD2R   1       0       0.0     0.0     0.0     0.0     0.0     1.0     1.0     0.0     0.0\n")
        
        # 添加材料属性
        f.write("MAT1     1       2.1+11          0.3\n")
        f.write("PSOLID   1       1\n")
        
        # 写入节点信息 (GRID cards)
        for i, point in enumerate(points, 1):
            # 使用更严格的格式控制，确保数值精度
            f.write(f"GRID*   {i:8d}        1               {point[0]:16.9f}{point[1]:16.9f}*\n")
            f.write(f"*       {point[2]:16.9f}\n")
        
        # 写入单元信息 (CTETRA/CHEXA cards)
        for i, (cell, cell_type) in enumerate(zip(cells, cell_types), 1):
            if cell_type == vtk.VTK_TETRA:
                # 四面体单元，添加属性ID
                f.write(f"CTETRA  {i:8d}1       {cell[0]+1:8d}{cell[1]+1:8d}{cell[2]+1:8d}{cell[3]+1:8d}\n")
            elif cell_type == vtk.VTK_HEXAHEDRON:
                # 六面体单元，添加属性ID
                f.write(f"CHEXA   {i:8d}1       {cell[0]+1:8d}{cell[1]+1:8d}{cell[2]+1:8d}{cell[3]+1:8d}+\n")
                f.write(f"+       {cell[4]+1:8d}{cell[5]+1:8d}{cell[6]+1:8d}{cell[7]+1:8d}\n")
        
        # 写入文件尾
        f.write("ENDDATA\n")
    
    print(f"转换完成! 已保存到: {nas_file}")

# 使用示例
if __name__ == "__main__":
    vtk_file = "finiteGrid.vtk"  # 输入的VTK文件路径
    nas_file = "finiteGrid.nas"  # 输出的NASTRAN文件路径
    vtk_to_nas(vtk_file, nas_file) 